// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#include <stdint.h>
#include <mutex>
#include <random>

#include <pollux/common/caching/cached_factory.h>
#include <pollux/common/caching/simple_lru_cache.h>
#include <pollux/common/io/io_statistics.h>

#pragma once

namespace kumo::pollux::dwio::common {
    /// A throttler that can be used to backoff IO when the storage is overloaded.
    class Throttler {
    public:
        /// The configuration of the throttler.
        struct Config {
            /// If true, enables throttling of IO.
            bool throttleEnabled;

            /// The minimum backoff duration in milliseconds.
            uint64_t minThrottleBackoffMs;

            /// The maximum backoff duration in milliseconds.
            uint64_t maxThrottleBackoffMs;

            /// The backoff duration scale factor.
            double backoffScaleFactor;

            /// The minimum number of received local throttled signals before starting
            /// backoff.
            uint32_t minLocalThrottledSignals;

            /// The minimum number of received global throttled signals before starting
            /// backoff.
            uint32_t minGlobalThrottledSignals;

            /// The minimum number of received network throttled signals before starting
            /// backoff.
            uint32_t minNetworkThrottledSignals;

            /// The maximum number of entries in the throttled signal cache. There is
            /// one cache for each throttle signal type. For local throttle signal
            /// cache, each cache entry corresponds to a unqiue file direcotry in a
            /// storage system. For global throttle signal cache, each entry corresponds
            /// to a unique storage system.
            uint32_t maxCacheEntries;

            /// The TTL of the throttled signal cache entries in milliseconds. We only
            /// track the recently throtted signals.
            uint32_t cacheTTLMs;

            static constexpr bool kThrottleEnabledDefault{true};
            static constexpr uint64_t kMinThrottleBackoffMsDefault{200};
            static constexpr uint64_t kMaxThrottleBackoffMsDefault{30'000};
            static constexpr double kBackoffScaleFactorDefault{2.0};
            static constexpr uint32_t kMinLocalThrottledSignalsDefault{1'000};
            static constexpr uint32_t kMinGlobalThrottledSignalsDefault{100'000};
            static constexpr uint32_t kMinNetworkThrottledSignal{1'000};
            static constexpr uint32_t kMaxCacheEntriesDefault{10'000};
            static constexpr uint32_t kCacheTTLMsDefault{3 * 60 * 1'000};

            Config(
                bool throttleEnabled = kThrottleEnabledDefault,
                uint64_t minThrottleBackoffMs = kMinThrottleBackoffMsDefault,
                uint64_t maxThrottleBackoffMs = kMaxThrottleBackoffMsDefault,
                double backoffScaleFactor = kBackoffScaleFactorDefault,
                uint32_t minLocalThrottledSignals = kMinLocalThrottledSignalsDefault,
                uint32_t minGlobalThrottledSignals = kMinGlobalThrottledSignalsDefault,
                uint32_t minNetworkThrottledSignals = kMinNetworkThrottledSignal,
                uint32_t maxCacheEntries = kMaxCacheEntriesDefault,
                uint32_t cacheTTLMs = kCacheTTLMsDefault);

            std::string toString() const;
        };

        /// The throttler stats.
        struct Stats {
            std::atomic_uint64_t localThrottled{0};
            std::atomic_uint64_t globalThrottled{0};
            std::atomic_uint64_t networkThrottled{0};
            /// Counts the backoff delay in milliseconds.
            io::IoCounter backOffDelay;
        };

        static void init(const Config &config);

        static Throttler *instance();

        /// The type of throttle signal type.
        enum class SignalType {
            /// No throttled signal.
            kNone,
            /// A file directory throttled signal.
            kLocal,
            /// A cluster-wise throttled signal.
            kGlobal,
            /// Network throttled signal.
            kNetwork,
        };

        static std::string signalTypeName(SignalType type);

        /// Invoked to backoff when received a throttled signal on a particular
        /// storage location. 'type' specifies the throttled signal type received from
        /// the storage system. 'cluster' specifies the storage system. A query system
        /// might access data from different storage systems. 'directory' specifies
        /// the file directory within the storage system. The function returns the
        /// actual throttled duration in milliseconds. It returns zero if not
        /// throttled.
        uint64_t throttleBackoff(
            SignalType type,
            const std::string &cluster,
            const std::string &directory);

        const Stats &stats() const {
            return stats_;
        }

        static void testingReset() {
            instanceRef().reset();
        }

    private:
        static melon::SharedMutex &instanceLock() {
            static melon::SharedMutex mu;
            return mu;
        }

        static std::unique_ptr<Throttler> &instanceRef() {
            static std::unique_ptr<Throttler> instance;
            return instance;
        }

        explicit Throttler(const Config &config);

        bool throttleEnabled() const {
            return throttleEnabled_;
        }

        // Calculates the delay in milliseconds with exponential backoff for a storage
        // location, using the signal counters in cache and flags in config, and
        // update the throttle signal caches.
        uint64_t calculateBackoffDurationAndUpdateThrottleCache(
            SignalType type,
            const std::string &cluster,
            const std::string &directory);

        struct ThrottleSignal {
            uint64_t count{0};

            explicit ThrottleSignal(uint64_t _count) : count(_count) {
            }
        };

        // Creates ThrottleSignal via the Generator interface the CachedFactory
        // requires.
        class ThrottleSignalGenerator {
        public:
            ThrottleSignalGenerator() = default;

            std::unique_ptr<ThrottleSignal> operator()(
                const std::string & /*unused*/,
                const void * /*unused*/,
                void * /*unused*/);
        };

        using CachedThrottleSignalPtr = CachedPtr<std::string, ThrottleSignal>;

        using ThrottleSignalFactory = kumo::pollux::
        CachedFactory<std::string, ThrottleSignal, ThrottleSignalGenerator>;

        struct ThrottleSignalCache {
            std::unique_ptr<ThrottleSignalFactory> throttleCache;
            uint32_t minThrottledSignalsToBackOff;
        };

        void updateThrottleCacheLocked(
            SignalType type,
            const std::string &cluster,
            const std::string &directory,
            CachedThrottleSignalPtr &localSignal,
            CachedThrottleSignalPtr &globalSignal,
            CachedThrottleSignalPtr &networkSignal);

        void updateThrottleStats(SignalType type, uint64_t backoffDelayMs);

        static ThrottleSignalCache maybeMakeThrottleSignalCache(
            bool enabled,
            uint32_t minThrottledSignals,
            uint32_t maxCacheEntries,
            uint32_t cacheTTLMs);

        static const uint64_t kNoBackOffMs_{0};

        const bool throttleEnabled_;
        const uint64_t minThrottleBackoffDurationMs_;
        const uint64_t maxThrottleBackoffDurationMs_;
        const double backoffScaleFactor_;

        const ThrottleSignalCache localThrottleCache_;
        const ThrottleSignalCache globalThrottleCache_;
        const ThrottleSignalCache networkThrottleCache_;

        mutable std::mutex mu_;

        std::mt19937 rng_;
        Stats stats_;
    };

    std::ostream &operator<<(std::ostream &os, Throttler::SignalType type);
} // namespace kumo::pollux::dwio::common
